package com.wutoon.etl.handler;

import cn.hutool.core.util.StrUtil;
import com.wutoon.etl.aop.PowerJobLogReport;
import com.wutoon.etl.constant.CommonConstant;
import com.wutoon.etl.pojo.DataTransferTask;
import com.wutoon.etl.util.CommonUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.sql.*;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicMarkableReference;

/**
 * @author 武兴云/72176468
 * @version 1.0
 * @date 2024/1/22 17:48
 * @description 读取数据
 */
@Slf4j
public class JdbcReader implements Runnable{
    private final BlockingQueue<Object[]> queue;

    private final ConcurrentHashMap<Integer, String> columnMap;

    private final DataTransferTask task;

    private final Connection conn;

    private final CountDownLatch countDownLatch;

    private final AtomicMarkableReference<String> msg;

    private RequestAttributes requestAttributes;

    public JdbcReader(BlockingQueue<Object[]> queue, ConcurrentHashMap<Integer, String> columnMap, DataTransferTask task, Connection conn,
                      CountDownLatch countDownLatch, AtomicMarkableReference<String> msg, RequestAttributes requestAttributes) {
        this.queue = queue;
        this.columnMap = columnMap;
        this.task = task;
        this.conn = conn;
        this.countDownLatch = countDownLatch;
        this.msg = msg;
        this.requestAttributes = requestAttributes;
    }

    @Override
    public void run() {
        // 子线程放入requestAttributes，方便上报日志
        RequestContextHolder.setRequestAttributes(requestAttributes);
        PowerJobLogReport.report("Jdbc reader start ......");
        log.info("Jdbc reader start ......");
        long dataSize = 0;
        long start = System.currentTimeMillis();
        ResultSet rs = null;
        /**
         * 多表join查询时prepareStatement不支持fetchSize
         */
        try (Statement stmt = conn.createStatement()) {
            stmt.setFetchSize(task.getFetchSize());
            stmt.setQueryTimeout(CommonConstant.SQL_QUERY_TIME_OUT);
            rs = stmt.executeQuery(task.getQuerySql());
            ResultSetMetaData metaData = rs.getMetaData();
            int columnCount = metaData.getColumnCount();
            List<Object[]> loseData = new ArrayList<>();
            for (int i = 0; i < columnCount; i++) {
                columnMap.putIfAbsent(i, StrUtil.toCamelCase(metaData.getColumnName(i + 1).toLowerCase(Locale.ROOT)));
            }
            // 有数据，并且主线程没有停止子线程
            while (rs.next() && !Thread.currentThread().isInterrupted()) {
                dataSize++;
                if (dataSize %  CommonConstant.MAX_FETCH_SIZE == 0) {
                    log.info("JdbcReader fetch data:{}", dataSize);
                    PowerJobLogReport.report("JdbcReader fetch data:" + dataSize);
                }
                Object[] row = new Object[columnCount];
                for (int i = 0; i < columnCount; i++) {
                    row[i] = rs.getObject(i + 1);
                }
                // 如果队列满了会阻塞等待一分钟
                boolean offer = queue.offer(row, CommonConstant.PRODUCER_TIMEOUT_SECONDS, TimeUnit.SECONDS);
                if (!offer) {
                    log.warn("JdbcReader offer time out:{}", row);
                    loseData.add(row);
                }
            }
            if (CollectionUtils.isNotEmpty(loseData)) {
                // 如果存在漏提数据，最后在提一次, 再次失败则认为任务失败
                boolean success;
                for(Object[] row : loseData) {
                    success = queue.offer(row, CommonConstant.PRODUCER_TIMEOUT_SECONDS, TimeUnit.SECONDS);
                    if (!success) {
                        throw new RuntimeException("DataTransfer reader lose data");
                    }
                }
            }
        } catch (Exception e) {
            log.error("DataTransfer reader error:, ", e);
            String stackTraceString = CommonUtil.getStackTraceString(e);
            msg.set(stackTraceString, true);
        } finally {
            countDownLatch.countDown();
            log.info("DataTransfer reader end duration:{}, dataSize:{}", System.currentTimeMillis() - start, dataSize);
            PowerJobLogReport.report("JdbcReader fetch duration:%s, dataSize:%s" , System.currentTimeMillis() - start, dataSize);
            if (rs != null) {
                try {
                    rs.close();
                } catch (SQLException e) {
                    log.error("JdbcReader rs close error", e);
                }
            }
            RequestContextHolder.resetRequestAttributes();
        }
    }

}
